package org.apache.flink.runtime.resourceprovider;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

import javax.annotation.Nullable;

import org.apache.flink.api.common.ConsumerID;
import org.apache.flink.runtime.clusterframework.types.ResourceID;
import org.apache.flink.runtime.providermanager.ProviderManagerActions;
import org.apache.flink.runtime.resourceconsumer.ResourceConsumerGateway;
import org.apache.flink.runtime.resourceconsumer.ConsumingMasterId;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.function.SupplierWithException;

/** Default implementation of the {@link ConsumerTable}. */
/**
 * @类名: DefaultJobTable
 * @说明: 
 *
 * @author   leehom
 * @Date	 2024年1月28日 下午9:59:34
 * 修改记录：
 *
 * @see 	 
 */
public final class DefaultConsumerTable implements ConsumerTable {
    private final Map<ConsumerID, ConsumerOrConnection> consumers;

    private final Map<ResourceID, ConsumerID> resourceIdJobIdIndex;

    private DefaultConsumerTable() {
        this.consumers = new HashMap<>();
        this.resourceIdJobIdIndex = new HashMap<>();
    }

    @Override
    public <E extends Exception> Consumer getOrCreateConsumer(
            ConsumerID consumerId,
            SupplierWithException<? extends ConsumerTable.ConsumerServices, E> jobServicesSupplier)
            throws E {
        ConsumerOrConnection job = consumers.get(consumerId);

        if (job == null) {
            job = new ConsumerOrConnection(consumerId, jobServicesSupplier.get());
            consumers.put(consumerId, job);
        }

        return job;
    }

    @Override
    public Optional<Consumer> getConsumer(ConsumerID jobId) {
        return Optional.ofNullable(consumers.get(jobId));
    }

    @Override
    public Optional<Connection> getConnection(ConsumerID jobId) {
        return getConsumer(jobId).flatMap(Consumer::asConnection);
    }

    @Override
    public Optional<Connection> getConnection(ResourceID resourceId) {
        final ConsumerID jobId = resourceIdJobIdIndex.get(resourceId);

        if (jobId != null) {
            return getConnection(jobId);
        } else {
            return Optional.empty();
        }
    }

    @Override
    public Collection<Consumer> getConsumers() {
        return new ArrayList<>(consumers.values());
    }

    @Override
    public boolean isEmpty() {
        return consumers.isEmpty();
    }

    public static DefaultConsumerTable create() {
        return new DefaultConsumerTable();
    }

    @Override
    public void close() {
        for (ConsumerTable.Consumer job : getConsumers()) {
            job.close();
        }
    }

    private final class ConsumerOrConnection implements ConsumerTable.Consumer, ConsumerTable.Connection {

        private final ConsumerID consumerId;

        private final ConsumerTable.ConsumerServices jobServices;

        @Nullable private EstablishedConnection connection;

        private boolean isClosed;

        private ConsumerOrConnection(ConsumerID jobId, ConsumerTable.ConsumerServices jobServices) {
            this.consumerId = jobId;
            this.jobServices = jobServices;
            this.connection = null;
            this.isClosed = false;
        }

        @Override
        public boolean isConnected() {
            verifyJobIsNotClosed();
            return connection != null;
        }

        @Override
        public ConsumerTable.Consumer disconnect() {
            resourceIdJobIdIndex.remove(verifyContainsEstablishedConnection().getResourceID());
            connection = null;

            return this;
        }

        @Override
        public ConsumingMasterId getConsumingMasterId() {
            return verifyContainsEstablishedConnection().getResourceConsumerId();
        }

        @Override
        public ResourceConsumerGateway getResourceConsumerGateway() {
            return verifyContainsEstablishedConnection().getResourceConsumerGateway();
        }

        @Override
        public ProviderManagerActions getTaskManagerActions() {
            return verifyContainsEstablishedConnection().getTaskManagerActions();
        }

        @Override
        public ConsumerID getConsumerId() {
            return consumerId;
        }

        @Override
        public ResourceID getResourceId() {
            return verifyContainsEstablishedConnection().getResourceID();
        }

        @Override
        public Optional<ConsumerTable.Connection> asConnection() {
            verifyJobIsNotClosed();
            if (connection != null) {
                return Optional.of(this);
            } else {
                return Optional.empty();
            }
        }

        @Override
        public ConsumerTable.Connection connect(
                ResourceID resourceId,
                ResourceConsumerGateway jobMasterGateway,
                ProviderManagerActions taskManagerActions) {
            verifyJobIsNotClosed();
            Preconditions.checkState(connection == null);

            connection =
                    new EstablishedConnection(
                            resourceId,
                            jobMasterGateway,
                            taskManagerActions);
            resourceIdJobIdIndex.put(resourceId, consumerId);

            return this;
        }

        @Override
        public void close() {
            if (!isClosed) {
                if (isConnected()) {
                    disconnect();
                }

                jobServices.close();
                consumers.remove(consumerId);

                isClosed = true;
            }
        }

        private void verifyJobIsNotClosed() {
            Preconditions.checkState(!isClosed, "The job has been closed.");
        }

        private EstablishedConnection verifyContainsEstablishedConnection() {
            verifyJobIsNotClosed();
            Preconditions.checkState(
                    connection != null, "The job has not been connected to a JobManager.");
            return connection;
        }
    }

    private static final class EstablishedConnection {

        // The unique id used for identifying the job manager
        private final ResourceID resourceID;

        // Gateway to the job master
        private final ResourceConsumerGateway resourceConsumerGateway;

        // Task manager actions with respect to the connected job manager
        private final ProviderManagerActions taskManagerActions;

        private EstablishedConnection(
                ResourceID resourceID,
                ResourceConsumerGateway jobMasterGateway,
                ProviderManagerActions taskManagerActions) {
            this.resourceID = Preconditions.checkNotNull(resourceID);
            this.resourceConsumerGateway = Preconditions.checkNotNull(jobMasterGateway);
            this.taskManagerActions = Preconditions.checkNotNull(taskManagerActions);
        }

        public ResourceID getResourceID() {
            return resourceID;
        }

        public ConsumingMasterId getResourceConsumerId() {
            return resourceConsumerGateway.getFencingToken();
        }

        public ResourceConsumerGateway getResourceConsumerGateway() {
            return resourceConsumerGateway;
        }

        public ProviderManagerActions getTaskManagerActions() {
            return taskManagerActions;
        }

    }
}
